/******************************************************************************
** This file is part of the cl-jpeg project.
**
** Copyright (C) 2011-2013 Vladimir Antonenko
**
** This program is free software; you can redistribute it and/or modify it
** under the terms of the GNU General Public License as published by the
** Free Software Foundation; either version 2 of the License,
** or (at your option) any later version.
**
** This program is distributed in the hope that it will be useful,
** but WITHOUT ANY WARRANTY; without even the implied warranty of
** MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
** See the GNU General Public License for more details.
**
** You should have received a copy of the GNU General Public License along
** with this program; if not, write to the Free Software Foundation, Inc.
******************************************************************************/

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <cl/opencl.h>
#include "wgetopt.h"
#include "jpegenc.h"

typedef unsigned char color;

#pragma once
#pragma pack(push)
#pragma pack(1)

typedef struct {
	color Blue;
	color Green;
	color Red;
} BGR;

#pragma pack(pop)

#ifdef CL_VERSION_1_2
cl_int clEnqueueMarker(cl_command_queue queue, cl_event *event)
{
	return clEnqueueMarkerWithWaitList(queue, 0, NULL, event);
}
#endif

static void print_usage(char* name) 
{
	printf("Usage: %s [options]\n", name);
	printf("Options:\n"
		"	-h - help\n"
		"	-P <n> - use platform #n\n"
		"	-D <n> - use device #n\n"
		"	-c - use CPU\n"
		"	-g - use GPU\n"
		"	-f <filename> - use kernel source\n"
		"	-k <kernel> - use kernel\n"
		"	-i <size> - work item size\n"
		"	-g <size> - work group size\n"
		"	-s <size> - image size WxH\n"
		"	-p - enable profiling\n"
		"	-t - check results\n"
		"	-v - be verbose\n"
	);
}

void PrintPlatformInfo(cl_platform_id id, unsigned verbose)
{
	cl_uint err;
	char    str[256];

	err = clGetPlatformInfo(id, CL_PLATFORM_NAME, sizeof(str), str, NULL);
	if (err == CL_SUCCESS) {
		printf("Platform: %s\n", str);
	}
	else fprintf(stderr, "Error: Failed to get platform info CL_PLATFORM_NAME!\n");

	err = clGetPlatformInfo(id, CL_PLATFORM_VERSION, sizeof(str), str, NULL);
	if (err == CL_SUCCESS) {
		printf("Version: %s\n", str);
	}
	else fprintf(stderr, "Error: Failed to get platform info CL_PLATFORM_VERSION!\n");

	if (verbose)
	{
		err = clGetPlatformInfo(id, CL_PLATFORM_EXTENSIONS, sizeof(str), str, NULL);
		if (err == CL_SUCCESS) {
			printf("Extensions: %s\n", str);
		}
		else fprintf(stderr, "Error: Failed to get platform info CL_PLATFORM_EXTENSIONS!\n");
	}
}

void PrintDeviceInfo(cl_device_id id, unsigned verbose)
{
	cl_uint        err;
	cl_ulong       size = 0;
	cl_device_type type;
	char           str[256];

	err = clGetDeviceInfo(id, CL_DEVICE_NAME, sizeof(str), str, NULL);
	if (err == CL_SUCCESS) {
		printf("Device: %s\n", str);
	}
	else fprintf(stderr, "Error: Failed to get device info CL_DEVICE_NAME!\n");

	err = clGetDeviceInfo(id, CL_DEVICE_VERSION, sizeof(str), str, NULL);
	if (err == CL_SUCCESS) {
		printf("Version: %s\n", str);
	}
	else fprintf(stderr, "Error: Failed to get device info CL_DEVICE_VERSION!\n");

	err = clGetDeviceInfo(id, CL_DEVICE_TYPE, sizeof(type), &type, NULL);
	if (err == CL_SUCCESS) {
		printf("Device Type: %d\n", type);
	}
	else fprintf(stderr, "Error: Failed to get device info CL_DEVICE_TYPE!\n");

	err = clGetDeviceInfo(id, CL_DEVICE_GLOBAL_MEM_SIZE, sizeof(size), &size, NULL);
	if (err == CL_SUCCESS) {
		printf("Global Mem Size: %lu\n", size);
	}
	else fprintf(stderr, "Error: Failed to get device info CL_DEVICE_GLOBAL_MEM_SIZE!\n");

	if (verbose)
	{
		err = clGetDeviceInfo(id, CL_DEVICE_EXTENSIONS, sizeof(str), str, NULL);
		if (err == CL_SUCCESS) {
			printf("Extensions: %s\n", str);
		}
		else fprintf(stderr, "Error: Failed to get device info CL_DEVICE_EXTENSIONS!\n");
	}
}

int GetOpenCLVersion(cl_device_id id)
{
	cl_uint err;
	int     major = 0;
	int	    minor = 0;
	char    str[256];

	err = clGetDeviceInfo(id, CL_DEVICE_OPENCL_C_VERSION, sizeof(str), str, NULL);
	if (err == CL_SUCCESS) {
		printf("%s\n", str);
		if (sscanf(str, "OpenCL C %d.%d", &major, &minor) == 2)
		{
			return major*10 + minor;
		}
	}

	return 0;
}

cl_int ReadKernelSource(const char *clname, char **source)
{
	size_t n;
	FILE   *fcl = fopen(clname, "r");

	if (!fcl) return CL_INVALID_PROGRAM;

	*source = malloc(65536);

	if (!*source) return CL_OUT_OF_HOST_MEMORY;

	n = fread(*source, 1, 65535, fcl);

	(*source)[n] = '\0';

	return CL_SUCCESS;
}


void rgb2yuv(unsigned i, float *data)
{
	float R, G, B;
	float Y, Cb, Cr;

	R = data[i+0];
	G = data[i+1];
	B = data[i+2];
	
	Y =          0.299f  * R + 0.587f  * G + 0.114f  * B;
	Cb = 128.f - 0.1687f * R - 0.3313f * G + 0.5f    * B;
	Cr = 128.f + 0.5f    * R - 0.4187f * G - 0.0813f * B;

	data[i+0] = Y;
	data[i+1] = Cb;
	data[i+2] = Cr;
}

static FILE *fjpg = NULL;
void write_jpeg(const unsigned char buff[], const unsigned size)
{
	fwrite(buff, size, 1, fjpg);
} 


#define SIZEX		sizex
#define SIZEY		sizey
#define IMAGE_SIZE	(SIZEX*SIZEY)

#define OPT_TEST_01		0x00000001
#define OPT_PROFILING	0x00010000

int main(int argc, char* argv[])
{
	int              devType = CL_DEVICE_TYPE_ALL;
	char             *clProgName = "jpeg.cl";
	char             *clKernelName = "trans";
	char             *source = NULL;

	cl_uint          nPlatform = 0;
	cl_uint          nDevice = 0;
	cl_platform_id   clPlatform[4];			// OpenCL platform list
	cl_uint          num_platforms;			// Number of platforms in the list
	cl_device_id     clDeviceID[8];			// Compute device id list
	cl_uint          num_devices;			// Number of devices in the list
	int              clVersion = 0;			// Device OpenCL version

	cl_context       clContext = NULL;		// compute clContext
	cl_command_queue clCommands = NULL;		// compute command queue
	cl_program       clProgram = NULL;		// compute clProgram
	cl_kernel        clKernel1 = NULL;		// compute clKernel #1
	cl_kernel        clKernel2 = NULL;		// compute clKernel #2
	cl_kernel        clKernel3 = NULL;		// compute clKernel #3
	cl_mem           clMemInput = NULL;		// device memory used for the clMemInput array
	cl_mem           clMemBuf0;				// intermeiate device memory
	cl_mem           clMemBuf1;				// intermeiate device memory
	cl_mem           clMemOutput;			// device memory used for the clMemOutput array
	cl_event         evStart;
	cl_event         evStop;
	cl_event         clEvent;
	BGR              *data = NULL;			// original data set given to device
	char             *result[8];			// results returned from device

	//size_t           offset[3];				// global domain offset
	size_t           global[3];				// global domain size
	size_t           local[3];				// local domain size
	size_t           item_size = 4;			// work item size
	size_t           group_size = 0;		// work group size
	cl_int           err;					// error code returned from api calls
	unsigned         verbose = 0;
	unsigned         options = 0;
	unsigned         count;
	//unsigned         i;
	int              sizex = 1024;
	int              sizey = 1024;
	cl_ulong         startTime;
	cl_ulong         endTime;
	cl_ulong         executionTime;
	char             opt;
	char             sProgramOpts[256];

	/* read command line */
	while ((opt = getopt(argc, argv, ":hP:D:cgf:k:i:G:s:ptv")) != -1)
	{
		switch(opt) {
		case 'h':
			print_usage(argv[0]);
			return 0;
		case 'P':
			nPlatform = atoi(optarg);
			break;
		case 'D':
			nDevice = atoi(optarg);
			break;
		case 'c':
			devType = CL_DEVICE_TYPE_CPU;
			break;
		case 'g':
			devType = CL_DEVICE_TYPE_GPU;
			break;
		case 'f':
			clProgName = optarg;
			break;
		case 'k':
			clKernelName = optarg;
			break;
		case 'i':
			item_size = atoi(optarg);
			break;
		case 'G':
			group_size = atoi(optarg);
			break;
		case 's':
			sscanf(optarg, "%dx%d", &sizex, &sizey);
			break;
		case 'p':
			options |= OPT_PROFILING;
			break;
		case 't':
			options |= OPT_TEST_01;
			break;
		case 'v':
			verbose = 1;
			break;
		case ':':
			fprintf(stderr, "Option -%c requires an operand\n", optopt);
			print_usage(argv[0]);
			return -1;
		case '?':
			fprintf(stderr, "Unrecognized option: -%c\n", optopt);
			print_usage(argv[0]);
			return -1;
		}
	}

	if (optind+2 > argc) {
		print_usage(argv[0]); return -1;
	}

	{
		char *ifilename = argv[optind];
		FILE *fi = fopen(ifilename, "rb");

		if (!fi) {
			fprintf(stderr, "Error: Cannot open file %s!\n", ifilename);
			return EXIT_FAILURE;
		}

		// create data for the run
		data = malloc(sizeof(BGR)*IMAGE_SIZE);    // original RGB data given to device
		fread(data, sizeof(BGR), IMAGE_SIZE, fi);
		fclose(fi);
	}

	// Connect to a compute device
	err = clGetPlatformIDs(4, clPlatform, &num_platforms);
	if (err != CL_SUCCESS) {
		fprintf(stderr, "Error: Failed to find a platform!\n");
		return EXIT_FAILURE;
	}

	printf("Platform: %dof%d\n", nPlatform+1, num_platforms);
	PrintPlatformInfo(clPlatform[nPlatform], verbose);

	if(devType == CL_DEVICE_TYPE_CPU)
		printf("Using: CL_DEVICE_TYPE_CPU\n");
	else if(devType == CL_DEVICE_TYPE_GPU)
		printf("Using: CL_DEVICE_TYPE_GPU\n");

	// Get a device of the appropriate type
	err = clGetDeviceIDs(clPlatform[nPlatform], devType, 8, clDeviceID, &num_devices);
	if (err != CL_SUCCESS) {
		fprintf(stderr, "Error: Failed to create a device group!\n");
		return EXIT_FAILURE;
	}

	printf("\nDevice: %dof%d\n", nDevice+1, num_devices);
	PrintDeviceInfo(clDeviceID[nDevice], verbose);
	clVersion = GetOpenCLVersion(clDeviceID[nDevice]);

	// Create a compute context
	clContext = clCreateContext(0, 1, clDeviceID+nDevice, NULL, NULL, &err);
	if (!clContext) {
		fprintf(stderr, "Error: Failed to create a compute clContext!\n");
		return EXIT_FAILURE;
	}

	err = ReadKernelSource(clProgName, &source);
	if (err != CL_SUCCESS) {
		fprintf(stderr, "Error: Cannot read clKernel source code!\n");
		return EXIT_FAILURE;
	}

	// Create the compute program from the source buffer
	clProgram = clCreateProgramWithSource(clContext, 1, (const char**) &source, NULL, &err);
	// free the buffer with a source code
	free(source);
	if (!clProgram) {
		fprintf(stderr, "Error: Failed to create a compute clProgram!\n");
		return EXIT_FAILURE;
	}

	printf("\nImage size = %dx%d\n", sizex, sizey);
	count = IMAGE_SIZE;

	printf("\nBuilding CL source..."); fflush(stdout);
	sprintf(sProgramOpts, "-cl-mad-enable -cl-fast-relaxed-math -DSIZEX=%d -DSIZEY=%d", SIZEX, SIZEY);
	// Build the program executable
	err = clBuildProgram(clProgram, 0, NULL, sProgramOpts, NULL, NULL);
	if (err != CL_SUCCESS)
	{
		size_t len;
		char   buffer[10000];

		fprintf(stderr, "Error: Failed to build clProgram executable!\n");

		clGetProgramBuildInfo(clProgram, clDeviceID[nDevice], CL_PROGRAM_BUILD_LOG,
			sizeof(buffer), buffer, &len);

		fprintf(stderr, "%s\n", buffer);
		exit(1);
	}

	// Create the compute kernels in the program
	clKernel1 = clCreateKernel(clProgram, "v4_rgb2yuv420", &err);
	if (!clKernel1 || err != CL_SUCCESS) {
		fprintf(stderr, "Error %d: Failed to create compute clKernel1!\n", err);
		exit(1);
	}

	clKernel2 = clCreateKernel(clProgram, "subsample3", &err);
	if (!clKernel2 || err != CL_SUCCESS) {
		fprintf(stderr, "Error %d: Failed to create compute clKernel2!\n", err);
		exit(1);
	}

	clKernel3 = clCreateKernel(clProgram, "dct3N", &err);
	if (!clKernel3 || err != CL_SUCCESS) {
		fprintf(stderr, "Error %d: Failed to create compute clKernel3!\n", err);
		exit(1);
	}
	printf("done\n");

	// Create the input device buffer
	clMemInput = clCreateBuffer(clContext, CL_MEM_READ_ONLY|CL_MEM_USE_HOST_PTR,
		sizeof(BGR)*count, data, NULL);
	// YUV 4:4:4 buffer
	clMemBuf0 = clCreateBuffer(clContext, CL_MEM_READ_WRITE, sizeof(char)*count*3, NULL, NULL);
	// UV 4:2:0 buffer
	clMemBuf1 = clCreateBuffer(clContext, CL_MEM_READ_WRITE, sizeof(char)*count/2, NULL, NULL);
	// Output DCT buffer
	clMemOutput = clCreateBuffer(clContext, CL_MEM_WRITE_ONLY, sizeof(short)*count*3/2, NULL, NULL);

	if (!clMemInput || !clMemBuf0 || !clMemBuf1 || !clMemOutput) {
		fprintf(stderr, "Error: Failed to allocate device memory!\n");
		exit(1);
	}   

	// Create a command queue
	clCommands = clCreateCommandQueue(clContext, clDeviceID[nDevice],
		(options & OPT_PROFILING)? CL_QUEUE_PROFILING_ENABLE:0, &err);
	if (!clCommands) {
		fprintf(stderr, "Error: Failed to create a command clCommands!\n");
		exit(1);
	}

	// Set the arguments to the compute kernel1 (rgb->yuv)
	err  = clSetKernelArg(clKernel1, 0, sizeof(cl_mem), &clMemInput);
	err |= clSetKernelArg(clKernel1, 1, sizeof(cl_mem), &clMemBuf0);
	if (err != CL_SUCCESS) {
		fprintf(stderr, "Error: Failed to set clKernel1 arguments! %d\n", err);
		exit(1);
	}

	// Set the arguments to the compute kernel2 (chroma subsample)
	err  = clSetKernelArg(clKernel2, 0, sizeof(cl_mem), &clMemBuf0);
	err |= clSetKernelArg(clKernel2, 1, sizeof(cl_mem), &clMemBuf1);
	if (err != CL_SUCCESS) {
		fprintf(stderr, "Error: Failed to set clKernel2 arguments! %d\n", err);
		exit(1);
	}

	// Set the arguments to the compute kernel3 (DCT)
	err  = clSetKernelArg(clKernel3, 0, sizeof(cl_mem), &clMemBuf0);
	err |= clSetKernelArg(clKernel3, 1, sizeof(cl_mem), &clMemOutput);
	if (err != CL_SUCCESS) {
		fprintf(stderr, "Error: Failed to set clKernel3 arguments! %d\n", err);
		exit(1);
	}

	if (group_size)
	{
		local[0] = group_size;
	}
	else
	{
		// Get the maximum work group size for executing the clKernel on the device
		err = clGetKernelWorkGroupInfo(clKernel1, clDeviceID[nDevice],
			CL_KERNEL_WORK_GROUP_SIZE, sizeof(local), local, NULL);
		if (err != CL_SUCCESS) {
			fprintf(stderr, "Error: Failed to retrieve kernel work group info! %d\n", err);
		}
	}
	printf("Local workgroup size = %d\n", local[0]);

	startTime = __rdtsc();
	clEnqueueMarker(clCommands, &evStart);

	/* Transfer the input vector into device memory
	err = clEnqueueWriteBuffer(clCommands, clMemInput, CL_FALSE, 0,
		sizeof(BGR)*count, data, 0, NULL, NULL);
	if (err != CL_SUCCESS) {
		fprintf(stderr, "Error: Failed to write to source array!\n");
		exit(1);
	}*/

	// rgb2yuv kernel
	global[0] = SIZEX/16;
	global[1] = SIZEY/2;

	err = clEnqueueNDRangeKernel(clCommands, clKernel1,
		2, NULL, global, NULL,
		0, NULL, &clEvent);

	if (err != CL_SUCCESS) {
		fprintf(stderr, "Error %d: Failed to execute kernel1!\n", err);
	}

	// DCT kernel - uniform Y,Cb,Cr processing
	global[0] = SIZEY*SIZEX/8*3/2; local[0] = 64;

	err = clEnqueueNDRangeKernel(clCommands, clKernel3,
		1, NULL, global, local,
		0, NULL, &clEvent);

	if (err != CL_SUCCESS) {
		fprintf(stderr, "Error %d: Failed to execute kernel3!\n", err);
	}

	/*err = clEnqueueReadBuffer(clCommands, clMemOutput[5],
		CL_FALSE, 0, sizeof(short)*count, result[5], 0, NULL, NULL);*/
	result[5] = clEnqueueMapBuffer(clCommands, clMemOutput,
		CL_FALSE, CL_MAP_READ, 0, sizeof(short)*count*3/2, 0, NULL, NULL, &err);

	if (err != CL_SUCCESS) {
		fprintf(stderr, "Error %d: Failed to read output array!\n", err);
	}

	clEnqueueMarker(clCommands, &evStop);

	printf("\nExecuting..."); fflush(stdout);
	// Wait for all commands to complete
	clFinish(clCommands);
	endTime = __rdtsc();
	printf("done\n");

	executionTime = endTime - startTime;
	printf("Execution time %u *1000 clk\n", executionTime/1000);

	if (options & OPT_PROFILING)
	{
		clGetEventProfilingInfo(evStart, CL_PROFILING_COMMAND_SUBMIT, sizeof(cl_ulong), &startTime, NULL);
		clGetEventProfilingInfo(evStop,  CL_PROFILING_COMMAND_SUBMIT, sizeof(cl_ulong), &endTime, NULL);
		executionTime = endTime - startTime;
		printf("Kernels execution time (profiler) %u mks\n", executionTime/1000);
	}

	if ((fjpg = fopen(argv[optind+1], "wb")))
	{
		short *Y  = (short*) result[5];
		short *Cb = Y + SIZEY*SIZEX;
		short *Cr = Cb + SIZEY*SIZEX/4;
		unsigned i,j;

		startTime = __rdtsc();
		huffman_start(SIZEY, SIZEX);

		// expect the DCT data in 8x8 arrays
		for (i = 0; i < SIZEY/16u; i++)
		for (j = 0; j < SIZEX/16u; j++)
		{
			huffman_encode(HUFFMAN_CTX_Y, Y + i*SIZEX*16 + j*2*64 + 0);
			huffman_encode(HUFFMAN_CTX_Y, Y + i*SIZEX*16 + j*2*64 + 64);
			huffman_encode(HUFFMAN_CTX_Y, Y + i*SIZEX*16 + (j+SIZEX/16)*2*64 + 0);
			huffman_encode(HUFFMAN_CTX_Y, Y + i*SIZEX*16 + (j+SIZEX/16)*2*64 + 64);
			huffman_encode(HUFFMAN_CTX_Cb, Cb + i*SIZEX*4 + j*64);
			huffman_encode(HUFFMAN_CTX_Cr, Cr + i*SIZEX*4 + j*64);
		}

		huffman_stop();
		endTime = __rdtsc();
		fclose(fjpg); 

		executionTime = endTime - startTime;
		printf("JPEG execution time %u *1000 clk\n", executionTime/1000);
	}
	else
		fprintf(stderr, "Error: cannot create %s\n", argv[optind+1]);

	// dump intermediate yuv data
	if (options & OPT_TEST_01)
	{
		FILE *fo = fopen("420.yuv", "wb");

		if (fo)
		{
			// get some intermediate data (depends on how the first kernel stores data)
			result[0] = clEnqueueMapBuffer(clCommands, clMemBuf0,
				CL_TRUE, CL_MAP_READ, 0, sizeof(char)*count*3/2, 0, NULL, NULL, &err);

			if (err == CL_SUCCESS)
			{
				fwrite(result[0], sizeof(char), count*3/2, fo);
				fclose(fo);
				clEnqueueUnmapMemObject(clCommands, clMemBuf0, result[0], 0, NULL, NULL);
			}
			else
				fprintf(stderr, "Error %d: Failed to read device buffer!\n", err);
		}
		else
			fprintf(stderr, "Error: Cannot create file %s!\n", "420.yuv");

		fo = fopen("420.dct", "wb");

		if (fo)
		{
			// store DCT data
			fwrite(result[5], sizeof(short), count*3/2, fo);
			fclose(fo);
		}
		else
			fprintf(stderr, "Error: Cannot create file %s!\n", "420.dct");
	}

	// Shutdown and cleanup
	clEnqueueUnmapMemObject(clCommands, clMemOutput, result[5], 0, NULL, NULL);

	free(data);

	clReleaseMemObject(clMemInput);
	clReleaseMemObject(clMemBuf0);
	clReleaseMemObject(clMemBuf1);
	clReleaseMemObject(clMemOutput);
	clReleaseProgram(clProgram);
	clReleaseKernel(clKernel1);
	clReleaseKernel(clKernel2);
	clReleaseKernel(clKernel3);
	clReleaseCommandQueue(clCommands);
	clReleaseContext(clContext);

	return 0;
}